from joblib import load
import numpy as np


class Classification(object):
    def __init__(self):
        self.knn = load('knn_classifien.joblib')
        self.scaler = load('feature_scaler.joblib')
        self.le = load('label_encoder.joblie')

    def get_predictor_category(self, new_data):
        # 转换为数组并保存特征顺序
        feature_order = ['eps', 'total_revenue_ps', 'undist_profit_ps', 'gross_margin', 'fcff', 'fcfe', 'tangible_asset',
                         'bps', 'grossprofit_margin', 'npta', 'roic']
        new_values = np.array([[new_data[col] for col in feature_order]])
        # 标准化数据
        new_scaled = self.scaler.transform(new_values)

        # 预测分类
        predicted_label = self.knn.predict(new_scaled)
        predicted_category = self.le.inverse_transform(predicted_label)
        return predicted_category[0]

if __name__ == '__main__':
    c1 = Classification()
    new_data = {
        'eps': "0.84",
        'total_revenue_ps': "16.8383",
        'undist_profit_ps': "8.3656",
        'gross_margin': "37901400000",
        'fcff': "154232000000",
        'fcfe': "195770000000",
        'tangible_asset': "1861070000000",
        'bps': "20.8979",
        'grossprofit_margin': "18.8665",
        'npta': "0.8815",
        'roic': "2.204"
    }
    predicted_category = c1.get_predictor_category(new_data)
    print(predicted_category)